{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import sisl\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In this analysis example, we will show how to plot the wavefunction for a periodic system (the same scheme may be used to plot molecular orbitals).\n",
    "\n",
    "The basic principle of plotting the real-space wave functions may be written as this (for a given $\\mathbf k$-point):\n",
    "\\begin{equation}\n",
    "   \\psi_i(\\mathbf k, \\mathbf r) = \\sum_\\nu e^{i\\mathbf k \\cdot \\mathbf r_\\nu}c_{i\\nu}\\phi_\\nu(\\mathbf r - \\mathbf r_\\nu),\n",
    "\\end{equation}\n",
    "where $c_{i\\nu}=\\langle\\tilde\\phi_\\nu|\\psi_i\\rangle$ is the coefficient for the $i$th eigenstate and $\\nu$ basis orbital and $\\phi_\\nu(\\mathbf r - \\mathbf r_\\nu)$ is the basis orbital $\\nu$ centred at $\\mathbf r_\\nu$.\n",
    "\n",
    "`sisl` will, in most cases, automatically read the necessary information from a Siesta run to be able to construct the $\\phi_\\nu$ basis functions in real-space. If the basis-information is not available `sisl` will inform you when trying to calculate real-space quantities.\n",
    "\n",
    "In this example, we will calculate the eigenstates for a given $\\mathbf k$-point for graphene and plot a cut through the real-space wavefunction at a fixed distance above the graphene plane."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Exercises\n",
    "\n",
    "1. Run Siesta.\n",
    "2. Read in the Hamiltonian using the `RUN.fdf` file, see e.g. [S 1](../S_01/run.ipynb).\n",
    "3. Calculate the eigenstate for the $\\Gamma$-point (see. `Hamiltonian.eigenstate`)\n",
    "4. Read the entry about the `EigenstateElectron` in the [sisl](https://zerothi.github.io/sisl/api/generated/sisl.physics.electron.EigenstateElectron.html).  \n",
    "   Figure out which method you should use in order to calculate the real-space wavefunction.  \n",
    "   Use the below method (`plot_grid`) to plot a cut through the wavefunction grid.  \n",
    "   *HINT*: this may be a useful grid `grid = Grid(0.05, lattice=H.geometry.lattice)`\n",
    "5. If you **don't** see a warning like this:\n",
    "\n",
    "        info:0: SislInfo: wavefunction: summing 18 different state coefficients, will continue silently!\n",
    "\n",
    "   then you have done it correctly! Look in the documentation and figure out how to take a *sub*set of the eigenstates and only plot a single one of them on the grid. This is important since otherwise you are plotting the super-position of all eigenstates at the $\\mathbf k$-point.  \n",
    "   *HINT*: If you have VMD/XCrySDen on your laptop it may be *fun* to plot the real-space quantities using cube files, if you want, figure out how to save the `Grid` into a cube/xsf file.\n",
    "6. Try and *play* with the supercell you pass to the `Grid` initialization and plot the wavefunction for different sizes of the grid.  \n",
    "    What do you see for increasing size of grids? Are there any periodicities, if so, what is the periodicity and why?\n",
    "7. Calculate the eigenstates such that the wavefunctions has a periodicity of $3$ along the first lattice vector and $2$ along the second lattice vector (only plot one of the eigenstates)  \n",
    "   *HINT*: $e^{i\\mathbf k \\cdot \\mathbf R}$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_grid(grid, plane_dist=1):\n",
    "    \"\"\" Plot the grid in either 1 (real-only) or 2 (complex wavefunction) graphs\n",
    "    \n",
    "    A cut through the grid will be plotted corresponding to `plane_dist` above the z-coordinate\n",
    "    \"\"\"\n",
    "    z_index = grid.index(plane_dist, 2)\n",
    "    x, y = np.mgrid[:grid.shape[0], :grid.shape[1]]\n",
    "    dcell = grid.dcell\n",
    "    x, y = x * dcell[0, 0] + y * dcell[1, 0], x * dcell[0, 1] + y * dcell[1, 1]\n",
    "    if grid.dtype in [np.complex64, np.complex128]:\n",
    "        fig, axs = plt.subplots(1, 2, figsize=(15, 5))\n",
    "        axs[0].contourf(x, y, grid.grid[:, :, z_index].real)\n",
    "        im = axs[1].contourf(x, y, grid.grid[:, :, z_index].imag)\n",
    "        for ax in axs:\n",
    "            ax.set_xlabel(r'$x$ [Ang]'); ax.set_ylabel(r'$y$ [Ang]')\n",
    "        axs[0].set_title('Real part')\n",
    "        axs[1].set_title('Imaginary part')\n",
    "    else:\n",
    "        fig, ax = plt.subplots(1, 1) ; axs = [ax]\n",
    "        axs = [ax]\n",
    "        im = ax.contourf(x, y, grid.grid[:, :, z_index])\n",
    "        ax.set_xlabel(r'$x$ [Ang]'); ax.set_ylabel(r'$y$ [Ang]')\n",
    "    # Also plot the atomic coordinates\n",
    "    try:\n",
    "        xyz = grid.geometry.xyz\n",
    "        for ax in axs:\n",
    "            ax.scatter(xyz[:, 0], xyz[:, 1], 50, 'k', alpha=.6)\n",
    "    except: pass\n",
    "    fig.colorbar(im);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Add code here to\n",
    "#   1) read in Hamiltonian with basis orbitals\n",
    "#   2) calculate the eigenstate for a given k-point\n",
    "#   3) create a grid to plot the wavefunction on\n",
    "#   4) plot the grid using the above method\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
